import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import models
import torchvision.transforms as transforms
from data_loader import MyDataset
# 定义UNet模型，使用ResNet-50作为骨干网络
from UNetResNet50 import UNetResNet50

# 超参数设置
epochs = 100
learning_rate = 1e-5
batch_size = 2**4  # 根据显存情况调整
num_classes = 1  # 假设是二分类问题
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据加载和预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 假设你有一个自定义的数据集类 MyDataset
train_dataset = MyDataset(transform=transform)
val_dataset = MyDataset(transform=transform)
test_dataset = MyDataset(transform=transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型、损失函数和优化器
import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet50",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=1,                      # model output channels (number of classes in your dataset)
)
criterion = nn.BCEWithLogitsLoss()  # 二分类问题
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练过程
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader)}")

    # 验证模型
    model.eval()
    with torch.no_grad():
        val_loss = 0.0
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item()

        print(f"Validation Loss: {val_loss/len(val_loader)}")

# 测试集上的评估
def evaluate_model(model, loader):
    model.eval()
    with torch.no_grad():
        metrics = {"PA": 0, "CPA": 0, "MPA": 0, "IoU": 0, "MIoU": 0}
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            # 在这里计算PA、CPA、MPA、IoU和MIoU等指标
            # 可以根据实际需要定义你自己的评估函数
        return metrics

test_metrics = evaluate_model(model, test_loader)
print(f"Test Metrics: {test_metrics}")
